package com.hujiang.basic.search.common.converter.req;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.hujiang.basic.search.common.converter.Converter;
import com.hujiang.basic.search.model.HJSearchRequest;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.elasticsearch.action.search.SearchRequest;

import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.range.date.DateRangeBuilder;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.highlight.HighlightBuilder;

import org.elasticsearch.search.rescore.RescoreBuilder;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.stereotype.Component;

@Component
public class SearchRequestConverter implements Converter<HJSearchRequest, SearchRequest> {

    @Override
    public SearchRequest convert(HJSearchRequest in) {
        SearchRequest request = new SearchRequest();
        request.indices(in.getIndex());
        request.types(in.getType());

        SearchSourceBuilder builder = new SearchSourceBuilder();

        if (MapUtils.isNotEmpty(in.getQuery())) {
            builder.query(in.getQuery());
        }

        if(MapUtils.isNotEmpty(in.getRescore())){
            buildRescore(in, builder);
        }

        if (CollectionUtils.isNotEmpty(in.getSort())) {
            buildSort(in, builder);
        }

        if (MapUtils.isNotEmpty(in.getAggs())) {
            buildAggs(in, builder);
        }

        if (CollectionUtils.isNotEmpty(in.getFields())) {
            builder.fields(in.getFields());
        }

        if (MapUtils.isNotEmpty(in.getHighlight())) {
            bulidHighlight(in, builder);
        }

        builder.from(in.getFrom());
        builder.size(in.getSize());

        request.source(builder);
        return request;
    }


    private void bulidHighlight(HJSearchRequest in, SearchSourceBuilder builder) {
        HighlightBuilder highlightBuilder = new HighlightBuilder();
        if (in.getHighlight().containsKey("pre_tags")) {
            JSONArray pre_tagsArray = (JSONArray) (in.getHighlight().get("pre_tags"));
            highlightBuilder.preTags(pre_tagsArray.toArray(new String[pre_tagsArray.size()]));
        }
        if (in.getHighlight().containsKey("post_tags")) {
            JSONArray post_tagsArray = (JSONArray) (in.getHighlight().get("post_tags"));
            highlightBuilder.postTags(post_tagsArray.toArray(new String[post_tagsArray.size()]));
        }
        ((JSONObject) in.getHighlight().get("fields")).forEach((k, v) -> {
            JSONObject fieldsValue = (JSONObject) v;
            if (fieldsValue.containsKey("fragment_size") && fieldsValue.containsKey("number_of_fragments")) {
                highlightBuilder.field(k, fieldsValue.getInteger("fragment_size"), fieldsValue.getInteger("number_of_fragments"));
            } else if(fieldsValue.containsKey("fragment_size")){
                highlightBuilder.field(k, fieldsValue.getInteger("fragment_size"));
            } else if(fieldsValue.containsKey("number_of_fragments")){
                highlightBuilder.field(k, 100, fieldsValue.getInteger("number_of_fragments"));
            }
            else {
                highlightBuilder.field(k);
            }
        });
        builder.highlight(highlightBuilder);
    }

    private void buildSort(HJSearchRequest in, SearchSourceBuilder builder) {
        for (Object sort : in.getSort()) {
            if ("_score".equals(sort)) {
                builder.sort(SortBuilders.scoreSort());
            } else {
                JSONObject sortFields = (JSONObject) sort;
                if (sortFields.size() > 0) {
                    sortFields.forEach((k, v) ->
                    {
                        if(v instanceof JSONObject && ((JSONObject)v).containsKey("order")){//适配原生sort
                            builder.sort(SortBuilders.fieldSort(k).order(SortOrder.valueOf(((JSONObject) v).getString("order").toUpperCase())));
                        }else{
                            builder.sort(SortBuilders.fieldSort(k).order(SortOrder.valueOf(v.toString().toUpperCase())));
                        }
                    });
                }
            }
        }
    }


//    "rescore": {
//                "window_size": 10,
//                "query": {
//                    "score_mode": "multiply",
//                    "query_weight": 0.7,
//                    "rescore_query_weight": 1.2,
//                    "rescore_query": {
//                       "function_score": {
//                          "script_score": {
//                             "script": {
//                                "inline": "doc['pc_lesson_exposure'].value"
//                        }
//                    }
//                }
//            }
//        }
//    }
    private void buildRescore(HJSearchRequest in, SearchSourceBuilder builder) {
        JSONObject rescoreJson = in.getRescore();
        RescoreBuilder rescoreBuilder = new RescoreBuilder();
        if(rescoreJson.containsKey("window_size")){
            int windowSize = rescoreJson.getInteger("window_size");
            rescoreBuilder.windowSize(windowSize);
        }
        if(rescoreJson.containsKey("query")){
            JSONObject queryJson = rescoreJson.getJSONObject("query");
            FunctionScoreQueryBuilder functionScoreQuery = null ;
            if(queryJson.containsKey("rescore_query")){
                JSONObject rescore_queryJson = queryJson.getJSONObject("rescore_query");
                if(rescore_queryJson.containsKey("function_score")){
                    JSONObject function_scoreJson = rescore_queryJson.getJSONObject("function_score");
                    if(function_scoreJson.containsKey("script_score")){
                        JSONObject script_scoreJson = function_scoreJson.getJSONObject("script_score");
                        if(script_scoreJson.containsKey("script")){
                            JSONObject scriptJson = script_scoreJson.getJSONObject("script");
                            if(scriptJson.containsKey("inline")){
                                 functionScoreQuery = QueryBuilders.functionScoreQuery(
                                         ScoreFunctionBuilders.scriptFunction(scriptJson.getString("inline")));
                            }
                        }
                    }
                }
            }
            if(null == functionScoreQuery) return ;
            RescoreBuilder.QueryRescorer queryRescorer = RescoreBuilder.queryRescorer(functionScoreQuery);
            if(queryJson.containsKey("score_mode")){
                queryRescorer.setScoreMode(queryJson.getString("score_mode"));
            }
            if(queryJson.containsKey("query_weight")){
                queryRescorer.setQueryWeight(queryJson.getFloat("query_weight"));
            }
            if(queryJson.containsKey("rescore_query_weight")){
                queryRescorer.setRescoreQueryWeight(queryJson.getFloat("rescore_query_weight"));
            }
            rescoreBuilder.rescorer(queryRescorer);
            builder.addRescorer(rescoreBuilder);
        }
    }

    private void buildAggs(HJSearchRequest in, SearchSourceBuilder builder) {
        in.getAggs().forEach((k, v) ->
                {
                    if (v instanceof JSONObject) {
                        JSONObject vJson =(JSONObject) v;
                        if (vJson.containsKey("terms")) {
                            JSONObject termJson = (JSONObject) (((JSONObject) v).get("terms"));
                            TermsBuilder termBuilder =
                                    AggregationBuilders.terms(k).field(termJson.getString("field"));
                            if (termJson.containsKey("size")) {
                                termBuilder.size(termJson.getIntValue("size"));
                            }
                            if (termJson.containsKey("order")
                                    && ((JSONObject) (termJson.get("order"))).containsKey("_count")) {
                                JSONObject orderJson = ((JSONObject) (termJson.get("order")));
                                termBuilder.order(Terms.Order.count(
                                        "asc".equalsIgnoreCase(orderJson.getString("_count"))));
                            }
                            builder.aggregation(termBuilder);
                        } else if(vJson.containsKey("cardinality")){
                            JSONObject cardinalityJson = (JSONObject) (vJson.get("cardinality"));
                            builder.aggregation(AggregationBuilders.cardinality(k).field(cardinalityJson.getString("field")));
                        }else if(vJson.containsKey("date_range")){
                            JSONObject rangeJson = (JSONObject) (vJson.get("date_range"));
                            DateRangeBuilder dateRangeBuilder = AggregationBuilders.dateRange(k).field(rangeJson.getString("field"));
                            JSONArray rangeArray = rangeJson.getJSONArray("ranges");
                            for (int i = 0; i < rangeArray.size(); i++) {
                                JSONObject range = rangeArray.getJSONObject(i);
                                if(range.containsKey("from") && range.containsKey("to") ){
                                    dateRangeBuilder.addRange(range.get("from"),range.get("to"));
                                }else if(range.containsKey("from")){
                                    dateRangeBuilder.addUnboundedFrom(range.get("from"));
                                }else if(range.containsKey("to")){
                                    dateRangeBuilder.addUnboundedTo(range.get("to"));
                                }
                            }
                            builder.aggregation(dateRangeBuilder);
                        }
                    }
                }
        );
    }
}


